from pydantic import BaseModel, Field
from langchain.tools import tool
from langchain.messages import SystemMessage

from model.Ark import ArkModel

def init_model_ep():
    ## 虽然 langchain 提供了 init_chat_model 的方法，但其实本质上，还是使用自定义初始化模型较多
    model = ArkModel().model
    print(model.invoke("请给我讲个关于程序员的笑话."))

def output_parser():
    ## 虽然 langchain 提供了 init_chat_model 的方法，但其实本质上，还是使用自定义初始化模型较多
    model = ArkModel().model

    class CustomerOutput(BaseModel):
        type: str = Field(description="笑话的类型")
        content: str = Field(description="笑话的内容")

    with_model = model.with_structured_output(CustomerOutput)
    print(with_model.invoke("请给我讲个关于程序员的笑话."))

def tools_model():
    @tool
    def get_weather(address: str) -> str:
        """ 传递一个地址，返回该地址的天气情况

            :Args
                address: 地址
         """
        print('工具被调用了...', address)
        return '晴天'

    model = ArkModel().model
    with_tool_model = model.bind_tools([get_weather])

    ## 调用invoke 之后，模型并不会理解调用工具，而是会返回一个 tool_calls 的调用列表，老版也是如此。
    response = with_tool_model.invoke("今天深圳的天气怎样？")

    ## 此时，需要手动去调用工具
    tool_calls = response.tool_calls
    result = []
    for tool_item in tool_calls:
        item_result = get_weather.invoke(tool_item)
        result.append(item_result)

    ## 如果想要完整的让大模型执行流程，此时需要拼接 messages 交给大模型执行
    base_messages = [{ 'role': 'user', 'content': '今天深圳的天气怎么样?' }]
    response_content = with_tool_model.invoke(base_messages +  result)
    print(response_content)

def run():
    tools_model()